import Data.Char (digitToInt, ord)
import Ascii
import BTree
import Data.Binary
import qualified Data.ByteString.Lazy as L

type Tree = BTree

sortRootTree :: [Tree] -> [Tree]
sortRootTree [] = []
sortRootTree (x:xs) = sortRootTree [x1 | x1<-xs, getNodeWeight x1 <= getNodeWeight x] ++ [x] ++ sortRootTree [x2 | x2<-xs, getNodeWeight x2 > getNodeWeight x]

-- auxilary function for merge Root Trees
resort :: Tree -> [Tree] -> [Tree]
resort newTree oldTrees = [x1 | x1<-oldTrees, getNodeWeight x1 <= getNodeWeight newTree] ++ [newTree] ++ [x2 | x2<-oldTrees, getNodeWeight x2 > getNodeWeight newTree]

-- generate huffman tree from sorted Root Trees
genHuffmanTree :: [Tree] -> Tree
genHuffmanTree (x1:x2:xs) = genHuffmanTree (resort (Node ((getNodeWeight x1 + getNodeWeight x2), '\x00', "") x1 x2) xs)
genHuffmanTree a = a!!0

-- info is depicited as "1a" "2b" and so on
createNodes :: [(Int, Char)] -> [Tree]
createNodes ((weight, ch):nodesInfo) = (Node (weight, ch, "") Empty Empty) : createNodes nodesInfo
createNodes [] = []

type HuffmanTable = [(Char, String)]
createHuffmanTable :: String -> HuffmanTable
createHuffmanTable = (walkHuffmanTree "") . genHuffmanTree . sortRootTree . createNodes . getAsciiWeightTuple

sortTupleList :: [(Char, String)] -> [(Char, String)]
sortTupleList ((c,str):xs) = sortTupleList [(x,str)| (x,str)<-xs, ord x <= ord c] ++ [(c,str)] ++ [(x,str)| (x,str)<-xs, ord x > ord c]
sortTupleList [] = []

getSortedHuffmanTable = sortTupleList . createHuffmanTable

findCode :: Char -> HuffmanTable -> String
-- tbl is dict sorted
findCode c ((c1,str):xs)
  | c == c1 = str
  | otherwise  = findCode c xs
findCode c [] = "" -- dropped if not matched

encodeFileWithTable :: String -> HuffmanTable -> String
encodeFileWithTable content tbl = case content of
                                    (x:xs) -> findCode x tbl++encodeFileWithTable xs tbl
                                    otherwise -> ""

binStrToInt :: String -> Word8
binStrToInt [] = 0
binStrToInt (c:cs) = fromIntegral (digitToInt c) + 2 * binStrToInt cs

getWord8FromBinStr :: String -> [Word8]
getWord8FromBinStr [] = []
getWord8FromBinStr str@(c:cs) = binStrToInt (take 8 str) : getWord8FromBinStr (drop 8 str)

{-compress text file-}
writeEncodeData :: FilePath -> IO Bool
writeEncodeData path = do
  content <- readFile path
  L.writeFile "result.bin" (L.pack $getWord8FromBinStr $encodeFileWithTable content $getSortedHuffmanTable  content)
  return True

main = do
  writeEncodeData "usb.c"
